# ---------------------------
# 1️⃣ 定义示例数据
# ---------------------------
import numpy as np
from langchain_core.example_selectors import SemanticSimilarityExampleSelector
from langchain_core.prompts import PromptTemplate, FewShotPromptTemplate
from langchain_openai import OpenAIEmbeddings
from langchain_community.vectorstores import FAISS


examples = [
    {
        "question": "这只基金今天是否适合买入?",
        "answer": "今天发布了 AI 相关的新闻报道，都是积极推进的买入。"
    },
    {
        "question": "风险是什么",
        "answer": "时长不稳定，美联储降息。"
    },
]

# ---------------------------
# 2️⃣ 定义每条示例的模板
# ---------------------------
example_prompt = PromptTemplate(
    input_variables=["question", "answer"],
    template="Question: {question}\nAnswer: {answer}\n"
)

# ---------------------------
# 3️⃣ 构建 FewShotPromptTemplate
# ---------------------------
few_shot_prompt = FewShotPromptTemplate(
    examples=examples,
    example_prompt=example_prompt,
    suffix="Question: {input}\nAnswer:",  # 用户输入问题的位置
    input_variables=["input"]
)

# ---------------------------
# 4️⃣ 生成最终 prompt
# ---------------------------
user_question = "大成科技混合C"
final_prompt = few_shot_prompt.format(input=user_question)

print(final_prompt)


# FewShotPromptTemplate.invoke 走的是 Runnable 通用实现 → 一定会触发 callback manager → 遇到 langchain.debug 就炸




# -------------------实列选择器
# 1) 定义示例
examples = [
    {"question": "基金今天适合买入吗？", "answer": "市场情绪较好，可以适量买入。"},
    {"question": "基金的风险是什么？", "answer": "风险主要是行业集中度高。"},
    {"question": "基金经理是谁？", "answer": "由张三担任基金经理。"},
    {"question": "基金的长期表现如何？", "answer": "近五年业绩持续跑赢同类。"},
]

# 2) 示例模板
example_prompt = PromptTemplate(
    input_variables=["question", "answer"],
    template="Q: {question}\nA: {answer}"
)

def local_embed(text):
    # 简单 demo，用字符长度作为向量，真实项目可用 TF-IDF、SentenceTransformer 等
    return np.array([len(text)])

# 生成所有示例的向量
example_vectors = [local_embed(ex["question"]) for ex in examples]


# 自定义 selector
class LocalSemanticSelector(SemanticSimilarityExampleSelector):
    def __init__(self, examples, example_prompt, vectors, k=2):
        super().__init__(examples, example_prompt)
        self.vectors = vectors
        self.k = k

    def select_examples(self, input_str):
        input_vec = local_embed(input_str)
        # 计算余弦相似度
        similarities = [np.dot(input_vec, v) / (np.linalg.norm(input_vec)*np.linalg.norm(v)) for v in self.vectors]
        # 选出 top k index
        topk_idx = np.argsort(similarities)[-self.k:][::-1]
        return [self.examples[i] for i in topk_idx]
    

selector = LocalSemanticSelector(examples, example_prompt, example_vectors, k=2)
selected = selector.select_examples("基金风险有哪些？")
for s in selected:
    print(s)


# 本地 embedding 函数
def local_embed(text):
    # 简单 demo，用字符长度作为向量，真实项目可用 TF-IDF、SentenceTransformer 等
    return np.array([len(text)])

# 生成所有示例的向量
example_vectors = [local_embed(ex["question"]) for ex in examples]

